"""模块导入"""
import numpy
import torch
import torch.nn as nn


"""定义网络"""
class oriNet(nn.Module):
    # 模型结构：四层全连接
    def __init__(self, input_size, hidden_size, output_size):
        super(oriNet, self).__init__()
        self.layer1 = nn.Linear(input_size, hidden_size)
        self.layer2 = nn.Linear(hidden_size, hidden_size)
        self.layer3 = nn.Linear(hidden_size, hidden_size)
        self.layer4 = nn.Linear(hidden_size, output_size)
    # 定义前向传播
    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        return x

class scatterNet(nn.Module):
    # 模型结构：四层全连接
    def __init__(self, input_size, hidden_size, output_size):
        super(scatterNet, self).__init__()
        self.net = torch.nn.Sequential(torch.nn.Linear(input_size, hidden_size),
                                       torch.nn.ReLU(),
                                       torch.nn.Linear(hidden_size, hidden_size),
                                       torch.nn.ReLU(),
                                       torch.nn.Linear(hidden_size, hidden_size),
                                       torch.nn.ReLU(),
                                       torch.nn.Linear(hidden_size, output_size))
    # 定义前向传播
    def forward(self, x):
        x = self.net(x)
        return x

def para_init(net):
    # 递归获得net的所有子代Module
    for op in net.modules():
        # 针对不同类型操作采用不同初始化方式
        if isinstance(op, nn.Linear):
            nn.init.normal_(op.weight.data, mean=0, std=0.1)
            nn.init.normal_(op.bias.data, mean=0, std=0.1)
        else:
            pass
